package com.xiaozhi.service;

import cn.hutool.core.collection.CollectionUtil;
import com.xiaozhi.entity.SearchResult;
import com.xiaozhi.utils.SiliconFlowEmbeddingUtil;
import com.xiaozhi.utils.SiliconFlowRerankUtil;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@Service
@Slf4j
public class VectorSearchService {
    private static final String COLLECTION_NAME = "knowledge_segments";
    private static final String ID_FIELD = "id";
    private static final String DOCUMENT_ID_FIELD = "document_id";
    private static final String SEGMENT_INDEX_FIELD = "segment_index";
    private static final String SEGMENT_ID_FIELD = "segment_id";
    private static final String VECTOR_FIELD = "vector";
    private static final String CONTENT_FIELD = "content";
    private static final String METADATA_FIELD = "metadata";
    private static final int VECTOR_DIMENSION = 1024; // 向量维度

    @Autowired
    SiliconFlowEmbeddingUtil siliconFlowEmbeddingUtil;

    @Autowired
    SiliconFlowRerankUtil siliconFlowRerankUtil;

    @Autowired
    private MilvusClientV2 milvusClient;
    // 在类的常量定义区域添加
    private static final float MIN_SCORE_THRESHOLD = 0.5f; //

    @PostConstruct
    public void init() {
        try {
            log.info("初始化 Milvus service........................");
            loadCollection();
            log.info("Milvus service 初始化成功........................");
        } catch (Exception e) {
            log.error("Milvus service initialization failed: " + e.getMessage(), e);
        }
    }

    /**
     * 加载集合到内存
     */
    public void loadCollection() {
        milvusClient.loadCollection(LoadCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .build());

        log.info("Milvus service 向量集合加载成功..........................................");
    }

    /**
     * 根据查询内容搜索相似的文档分段（带重排），可限定文档范围
     *
     * @param query       查询内容
     * @param documentIds 限定的文档ID列表，如果为空则不限制
     * @param topK        返回最相似的K个结果
     * @return 相似的文档分段列表
     */
    public List<SearchResult> searchWithRerank(String query, List<String> documentIds, int topK) {
        try {
            if (CollectionUtil.isEmpty(documentIds)) {
                return new ArrayList<>();

            }
            // 生成查询向量
            List<Float> queryVector = siliconFlowEmbeddingUtil.generateEmbedding(query);
            // 在Milvus中搜索相似向量，可限定文档范围
            List<SearchResult> initialResults = this.searchSimilarVectorsWithContext(queryVector, documentIds, topK * 2); // 获取更多候选结果

            // 按文档ID和片段索引排序，便于后续合并
            initialResults.sort(Comparator.comparing(SearchResult::getDocumentId)
                    .thenComparingInt(SearchResult::getSegmentIndex));

            // 合并相邻片段
            List<SearchResult> mergedResults = mergeAdjacentSegments(initialResults);

            // 提取内容用于重排
            List<String> documents = mergedResults.stream()
                    .map(SearchResult::getContent)
                    .collect(Collectors.toList());
            if (documents.isEmpty()) {
                return initialResults;
            }

            // 使用重排模型对结果进行重新排序
            List<SiliconFlowRerankUtil.RerankResponse.Result> rerankedResults =
                    siliconFlowRerankUtil.rerank(query, documents, topK);
            // 根据重排结果重新排序初始结果
            List<SearchResult> finalResults = new ArrayList<>();
            if (!CollectionUtil.isEmpty(rerankedResults)) {
                SiliconFlowRerankUtil.RerankResponse.Result result = rerankedResults.get(0);
                double relevanceScore = result.getRelevance_score();
               if (relevanceScore > 0.3) {
                   SiliconFlowRerankUtil.RerankResponse.Document document = result.getDocument();
                   SearchResult originalResult = new SearchResult();
                   originalResult.setContent(document.getText());
                    originalResult.setScore((float) relevanceScore);
                    finalResults.add(originalResult);
               }
            }
            return finalResults;
        } catch (Exception e) {
            log.error("重排搜索失败: " + e.getMessage(), e);
            // 如果重排失败，返回原始搜索结果
            return search(query, documentIds, topK);
        }
    }

    /**
     * 根据查询内容搜索相似的文档分段，可限定文档范围
     *
     * @param query       查询内容
     * @param documentIds 限定的文档ID列表，如果为空则不限制
     * @param topK        返回最相似的K个结果
     * @return 相似的文档分段列表
     */
    public List<SearchResult> search(String query, List<String> documentIds, int topK) {
        try {
            // 生成查询向量
            List<Float> queryVector = siliconFlowEmbeddingUtil.generateEmbedding(query);
            // 在Milvus中搜索相似向量，可限定文档范围
            return this.searchSimilarVectors(queryVector, documentIds, topK);
        } catch (Exception e) {
            log.error("向量搜索失败: " + e.getMessage(), e);
            // 返回空列表而不是抛出异常，提高系统稳定性
            return new ArrayList<>();
        }
    }


    /**
     * 搜索相似向量，可限定文档范围
     */
    public List<SearchResult> searchSimilarVectors(List<Float> queryVector, List<String> documentIds, int topK) {
        // 转换向量格式
        float[] floatArray = new float[queryVector.size()];
        for (int i = 0; i < queryVector.size(); i++) {
            floatArray[i] = queryVector.get(i);
        }
        // 创建搜索向量
        FloatVec floatVec = new FloatVec(floatArray);

        // 构造搜索请求
        SearchReq.SearchReqBuilder searchReqBuilder = SearchReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(Collections.singletonList(floatVec))
                .annsField(VECTOR_FIELD)
                .topK(topK)
                .outputFields(Arrays.asList(ID_FIELD, DOCUMENT_ID_FIELD, SEGMENT_INDEX_FIELD, SEGMENT_ID_FIELD,CONTENT_FIELD, METADATA_FIELD));

        // 如果指定了documentIds，则添加过滤条件
        if (documentIds != null && !documentIds.isEmpty()) {
            StringBuilder filterBuilder = new StringBuilder();
            filterBuilder.append(DOCUMENT_ID_FIELD).append(" in [");
            for (int i = 0; i < documentIds.size(); i++) {
                if (i > 0) {
                    filterBuilder.append(",");
                }
                filterBuilder.append("'").append(documentIds.get(i)).append("'");
            }
            filterBuilder.append("]");
            searchReqBuilder.filter(filterBuilder.toString());
        }

        // 执行搜索
        SearchResp response = milvusClient.search(searchReqBuilder.build());
        // 处理搜索结果
        List<SearchResult> results = new ArrayList<>();
        List<List<io.milvus.v2.service.vector.response.SearchResp.SearchResult>> searchResults = response.getSearchResults();
        for (List<io.milvus.v2.service.vector.response.SearchResp.SearchResult> resultPerVector : searchResults) {
            for (io.milvus.v2.service.vector.response.SearchResp.SearchResult item : resultPerVector) {
                // 添加分值过滤，只返回高于最低阈值的结果
                if (item.getScore() >= MIN_SCORE_THRESHOLD) {
                    SearchResult result = new SearchResult();
                    result.setId((String) item.getEntity().get(ID_FIELD));
                    result.setDocumentId((String) item.getEntity().get(DOCUMENT_ID_FIELD));
                    result.setSegmentIndex(((Long) item.getEntity().get(SEGMENT_INDEX_FIELD)).intValue());
                    result.setContent((String) item.getEntity().get(CONTENT_FIELD));
                    result.setSegmentId((String) item.getEntity().get(SEGMENT_ID_FIELD));
                    result.setScore(item.getScore());
                    result.setMetadata((String) item.getEntity().get(METADATA_FIELD));
                    results.add(result);
                }
            }
        }
        return results;
    }


    /**
     * 搜索相似向量并扩展上下文，可限定文档范围
     */
    private List<SearchResult> searchSimilarVectorsWithContext(List<Float> queryVector, List<String> documentIds, int topK) {
        // 先获取基础的匹配结果
        List<SearchResult> baseResults = searchSimilarVectors(queryVector, documentIds, topK);
        System.out.println("基础匹配结果："+baseResults);

        // 为每个匹配结果获取相邻片段
        Set<String> retrievedSegmentIds = baseResults.stream()
                .map(SearchResult::getSegmentId)
                .collect(Collectors.toSet());

        List<SearchResult> extendedResults = new ArrayList<>(baseResults);

        for (SearchResult result : baseResults) {
            // 获取前一个片段
            SearchResult previousSegment = getAdjacentSegment(result.getDocumentId(), result.getSegmentIndex() - 1);
            if (previousSegment != null && !retrievedSegmentIds.contains(previousSegment.getSegmentId())) {
                extendedResults.add(previousSegment);
                retrievedSegmentIds.add(previousSegment.getSegmentId());
            }

            // 获取后一个片段
            SearchResult nextSegment = getAdjacentSegment(result.getDocumentId(), result.getSegmentIndex() + 1);
            if (nextSegment != null && !retrievedSegmentIds.contains(nextSegment.getSegmentId())) {
                extendedResults.add(nextSegment);
                retrievedSegmentIds.add(nextSegment.getSegmentId());
            }
        }
        System.out.println("处理后的结果："+baseResults);
        return extendedResults;
    }


    /**
     * 合并相邻的文档片段
     */
    private List<SearchResult> mergeAdjacentSegments(List<SearchResult> segments) {
        // 按文档ID分组
        Map<String, List<SearchResult>> groupedByDocument = segments.stream()
                .collect(Collectors.groupingBy(SearchResult::getDocumentId));

        List<SearchResult> mergedResults = new ArrayList<>();

        // 遍历每个文档的片段
        for (List<SearchResult> docSegments : groupedByDocument.values()) {
            // 按片段索引排序
            docSegments.sort(Comparator.comparingInt(SearchResult::getSegmentIndex));

            int i = 0;
            while (i < docSegments.size()) {
                SearchResult current = docSegments.get(i);
                // 创建新的合并片段
                SearchResult mergedSegment = new SearchResult();
                mergedSegment.setId(current.getId());
                mergedSegment.setDocumentId(current.getDocumentId());
                mergedSegment.setSegmentIndex(current.getSegmentIndex());
                mergedSegment.setSegmentId(current.getSegmentId());

                // 初始化内容构建器
                StringBuilder contentBuilder = new StringBuilder(current.getContent());
                float maxScore = current.getScore();

                // 查找连续的片段并合并
                int j = i + 1;
                while (j < docSegments.size() &&
                        docSegments.get(j).getSegmentIndex() == docSegments.get(j-1).getSegmentIndex() + 1) {
                    contentBuilder.append("\n").append(docSegments.get(j).getContent());
                    maxScore = Math.max(maxScore, docSegments.get(j).getScore());
                    j++;
                }

                // 设置合并后的内容和最高分数
                mergedSegment.setContent(contentBuilder.toString());
                mergedSegment.setScore(maxScore);
                mergedSegment.setMetadata(current.getMetadata());

                mergedResults.add(mergedSegment);
                i = j;
            }
        }

        return mergedResults;
    }

    /**
     * 根据文档ID和片段索引获取特定片段
     */
    private SearchResult getAdjacentSegment(String documentId, int segmentIndex) {
        try {
            // 构造查询条件
            String filter = DOCUMENT_ID_FIELD + " == '" + documentId + "' and " + SEGMENT_INDEX_FIELD + " == " + segmentIndex;

            // 创建占位符向量
            float[] placeholderVector = new float[VECTOR_DIMENSION];

            // 创建搜索请求
            SearchReq searchReq = SearchReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .data(Collections.singletonList(new FloatVec(placeholderVector)))
                    .annsField(VECTOR_FIELD)
                    .topK(1)
                    .filter(filter)
                    .outputFields(Arrays.asList(
                            ID_FIELD,
                            DOCUMENT_ID_FIELD,
                            SEGMENT_INDEX_FIELD,
                            SEGMENT_ID_FIELD,
                            CONTENT_FIELD,
                            METADATA_FIELD))
                    .build();

            // 执行搜索
            SearchResp response = milvusClient.search(searchReq);

            // 处理结果
            List<List<io.milvus.v2.service.vector.response.SearchResp.SearchResult>> searchResults = response.getSearchResults();
            if (!searchResults.isEmpty() && !searchResults.get(0).isEmpty()) {
                io.milvus.v2.service.vector.response.SearchResp.SearchResult item = searchResults.get(0).get(0);

                SearchResult result = new SearchResult();
                result.setId((String) item.getEntity().get(ID_FIELD));
                result.setDocumentId((String) item.getEntity().get(DOCUMENT_ID_FIELD));
                result.setSegmentIndex(((Long) item.getEntity().get(SEGMENT_INDEX_FIELD)).intValue());
                result.setContent((String) item.getEntity().get(CONTENT_FIELD));
                result.setSegmentId((String) item.getEntity().get(SEGMENT_ID_FIELD));
                result.setScore(item.getScore());
                result.setMetadata((String) item.getEntity().get(METADATA_FIELD));

                return result;
            }
        } catch (Exception e) {
            log.warn("获取相邻片段失败: documentId={}, segmentIndex={}", documentId, segmentIndex, e);
        }

        return null;
    }


}
